package MDBAWS;

# Stop us from shooting outselves in the foot
use strict;
use warnings;

# We are storing the configuration file in INI format
use File::Temp qw( tempdir );
use File::Basename;

use Data::Dumper;

# We need to interact with Amazon S3
use Net::Amazon::S3;
use Net::Amazon::S3::Bucket;

# We need to interact with Amazon SimpleDB
use Amazon::SimpleDB::Client;

# We need to interact with Amazon SQS
use Amazon::SQS::Simple;

use MDBUtils;

sub store_sample {

=head1 store_sample( $filename, [$metadata] )

Uploads sample to Amazon S3 and records the metadata to Amazon SimpleDB

I<Arguments>

=over

=item $filename

Full path to filename

=item $metadata

hashref of data to include in the SimpleDB upload

=back

I<Returns>

Nothing

I<Exceptions>

problems with the distribution extraction, write errors on the file system, ...

=cut

    my $filename = shift
      or die("FUNCTION USAGE: store_sample( filename, [metadata] );\n");
    my $metadata;

    $metadata = shift;

    print "MDBAWS::store_sample()\n";

    my $config = MDBUtils::get_config();

    if ( !defined( $metadata->{'rcvd_DateTime'} ) ) {
        $metadata->{'rcvd_DateTime'} = DateTime->now();
        $metadata->{'rcvd_DateTime'}->set_time_zone("GMT");
    }

    if ( !defined( $metadata->{'submitter'} ) ) {
        $metadata->{'submitter'} = $config->{'default'}->{'submitter'};
    }

    $metadata->{'original_filename'} = basename($filename);
    my $dir = tempdir( CLEANUP => 1 );

    ########################
    # Checksum the payload #
    ########################

    $metadata->{'sha256'} = MDBUtils::mdb_sha256_hex($filename);
    $metadata->{'sha1'}   = MDBUtils::mdb_sha1_hex($filename);
    $metadata->{'md5'}    = MDBUtils::mdb_md5_hex($filename);
    $metadata->{'size'}   = -s $filename;

    ########################
    # Create a ZIP archive #
    ########################

    # Define the filename of the ZIP archive
    $metadata->{'zipname'} = $metadata->{'sha256'} . ".zip";

    MDBUtils::compress_zip( $filename, $dir . $metadata->{'zipname'},
        "infected" );

    ################################
    # Upload the ZIP archive to S3 #
    ################################

    s3_put_file(
        $config->{'aws'}->{'bucketname'},
        $dir . $metadata->{'zipname'},
        $metadata->{'zipname'}
    );

    $metadata->{'sample_url'} =
        "http://"
      . $config->{'aws'}->{'bucketname'} . "/"
      . $metadata->{'zipname'};

    # Delete the temp ZIP file
    unlink( $dir . $metadata->{'zipname'} );

    ###############################
    # Upload metadata to SimpleDB #
    ###############################

    sdb_put_metadata($metadata);

    ########################
    # Create SQS Ticket(s) #
    ########################

    sqs_put_queue( "store_sample", $metadata->{'sha256'} );
}

sub s3_put_file {

=head1 s3_put_file( $bucketname, $file, $filename )

Uploads $file as $filname to S3 bucket $bucketname

I<Arguments>

=over

=item $bucketname

Name of S3 bucket

=item $file

Full path to file to upload

=item $filename

Name of uploaded file

=back

I<Returns>

Nothing

I<Exceptions>

problems with the distribution extraction, write errors on the file system, ...

=cut

    my ( $bucketname, $file, $filename ) = @_;

    print "MDBAWS::s3_put_file()\n";

    if ( !$filename ) {
        $filename = $file;
    }

    my $config = MDBUtils::get_config();

    my $s3 = Net::Amazon::S3->new(
        aws_access_key_id     => $config->{'aws'}->{'aws_access_key_id'},
        aws_secret_access_key => $config->{'aws'}->{'aws_secret_access_key'},
        retry                 => 1
    );

    my $bucket = $s3->bucket($bucketname);
    my $ok = $bucket->add_key_filename( $filename, $file );

    if ( !$ok ) {
        print "ERROR: s3->add_key_filename( ) failed: "
          . $bucket->errstr() . "\n";
    }
}

sub s3_get_file {

=head1 $file = s3_get_file( $bucketname, $filename )

Downloads $filename from S3 bucket $bucketname. Returns path to downloaded file.

I<Arguments>

=over

=item $bucketname

Name of S3 bucket

=item $filename

Name of file to download

=back

I<Returns>

Path to downloaded file

I<Exceptions>

problems with the distribution extraction, write errors on the file system, ...

=cut

    my ( $bucketname, $filename ) = @_;
    my $config         = MDBUtils::get_config();
    my $dir            = $config->{'aws'}->{'s3_cache_dir'};
    my $local_filename = $dir . "/" . $filename;

    print "MDBAWS::s3_get_file()\n";

    if ( -e $local_filename ) {
        return ($local_filename);
    }

    my $s3 = Net::Amazon::S3->new(
        aws_access_key_id     => $config->{'aws'}->{'aws_access_key_id'},
        aws_secret_access_key => $config->{'aws'}->{'aws_secret_access_key'},
        retry                 => 1
    );

    my $bucket = $s3->bucket($bucketname);
    my $ok = $bucket->get_key_filename( $filename, "GET", $local_filename );

    if ( !$ok ) {
        print "ERROR: s3->get_key_filename() failed: "
          . $bucket->errstr() . "\n";
    }
    else {
        return ($local_filename);
    }
}

sub sdb_put_metadata {

=head1 sdb_put_metadata( $metadata )

Stores metadata in SimpleDB

I<Arguments>

=over

=item $metadata

Hash reference containing the metadata that whiches to be stored

=back

I<Returns>

Nothing

I<Exceptions>

problems with the distribution extraction, write errors on the file system, ...

=cut

    print "MDBAWS::sdb_put_metadata()\n";

    my $metadata = shift or die("");

    my $config = MDBUtils::get_config();

    my $simpledb = Amazon::SimpleDB::Client->new(
        $config->{'aws'}->{'aws_access_key_id'},
        $config->{'aws'}->{'aws_secret_access_key'}
    );

    my @attributes = ();

    # Put the metadata hash into an array, just like how SimpleDB wants it
    while ( my ( $key, $value ) = each(%$metadata) ) {

        if ( defined($key) && defined($value) ) {

            @attributes = {
                Name    => $key,
                Value   => $value,
                Replace => '1'
            };

            my $response = $simpledb->putAttributes(
                {
                    DomainName => $config->{'aws'}->{'simpledb_name'},
                    ItemName   => $metadata->{'sha256'},
                    Attribute  => @attributes
                }
            );
        }
    }
}

sub sdb_get_metadata {

=head1 $metadata = sdb_get_metadata( $item )

Get metadata from SimpleDB

I<Arguments>

=over

=item $item

Item to get

=back

I<Returns>

=over

=item $metadata

Hash reference containing the metadata stored in SimpleDB

=back

I<Exceptions>

problems with the distribution extraction, write errors on the file system, ...

=cut

    my $item = shift or die;

    print "MDBAWS::sdb_get_metadata()\n";

    my $config = MDBUtils::get_config();
    my $metadata;

    my $simpledb = Amazon::SimpleDB::Client->new(
        $config->{'aws'}->{'aws_access_key_id'},
        $config->{'aws'}->{'aws_secret_access_key'}
    );

    my $response = $simpledb->getAttributes(
        {
            DomainName => $config->{'aws'}->{'simpledb_name'},
            ItemName   => $item
        }
    );

    # Feed the SimpleDB data into an hash
    if ( $response->isSetGetAttributesResult() ) {
        my $getAttributesResult = $response->getGetAttributesResult();
        my $attributeList       = $getAttributesResult->getAttribute();
        foreach (@$attributeList) {
            my $attribute = $_;
            $metadata->{ $attribute->getName() } = $attribute->getValue();
        }
        return ($metadata);
    }
    else {
        return (0);
    }
}

sub sqs_put_queue {

=head1 sqs_put_queue( $modulename, $message )

Put $message in all output queues for $modulename

I<Arguments>

=over

=item $modulename

Name of module to output data as

=item $message

Message to send

=back

I<Returns>

Nothing

=back

I<Exceptions>

problems with the distribution extraction, write errors on the file system, ...

=cut

    my ( $modulename, $message ) = @_ or die;

    print "MDBAWS::sqs_put_queue()\n";

    my $config = MDBUtils::get_config();

    my $sqs = new Amazon::SQS::Simple(
        $config->{'aws'}->{'aws_access_key_id'},
        $config->{'aws'}->{'aws_secret_access_key'}
    );

    my (@out_queues) = split( /,/, $config->{$modulename}->{'out_queue'} );
    my $out_queue;
    foreach $out_queue (@out_queues) {
        my $queuename = $config->{'aws'}->{'sqs_prefix'} . $out_queue;
        my $q         = $sqs->CreateQueue($queuename);
        $q->SendMessage($message);
    }
}

sub sqs_get_queue {

=head1 $queue = sqs_get_queue( $modulename )

Receives $message from $modulename queue

I<Arguments>

=over

=item $modulename

Name of module to input data as

=back

I<Returns>

=over

=item $queue

SQS queue object

=back

I<Exceptions>

problems with the distribution extraction, write errors on the file system, ...

=cut

    my $modulename = shift or die;

    print "MDBAWS::sqs_get_queue()\n";

    my $config = MDBUtils::get_config();

    my $sqs = new Amazon::SQS::Simple(
        $config->{'aws'}->{'aws_access_key_id'},
        $config->{'aws'}->{'aws_secret_access_key'}
    );

    my $queue =
      $sqs->GetQueue( $config->{'aws'}->{'sqs_url'}
          . $config->{'aws'}->{'sqs_prefix'}
          . $config->{$modulename}->{'in_queue'} );

    return ($queue);
}

sub sqs_get_message {

=head1 $msg = sqs_get_message( $queue )

Receives $message from $modulename queue

I<Arguments>

=over

=item $queue

SQS Queue structure from sqs_get_queue()

=back

I<Returns>

=over

=item $message

SQS Message object

=back

I<Exceptions>

problems with the distribution extraction, write errors on the file system, ...

=cut

    my $queue = shift or die;

    print "MDBAWS::sqs_get_message()\n";

    my $msg = $queue->ReceiveMessage();

    if ( defined($msg)
        && $msg->MessageBody() =~
        m/e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855/ )
    {
        $queue->DeleteMessage( $msg->ReceiptHandle() );
        return ( sqs_get_message($queue) );
    }

    return ($msg);
}

1;
